context("Mutate")

test_that("mutate_epi.netsim", {
  nw <- network.initialize(n = 100, bipartite = 50, directed = FALSE)
  formation <- ~edges
  target.stats <- 50
  coef.diss <- dissolution_coefs(dissolution = ~offset(edges), duration = 20)
  est1 <- netest(nw, formation, target.stats, coef.diss, verbose = FALSE)

  # Epidemic model
  param <- param.net(inf.prob = 0.3, inf.prob.m2 = 0.15)
  init <- init.net(i.num = 1, i.num.m2 = 0)
  control <- control.net(type = "SI", nsteps = 10, nsims = 2,
                         verbose = FALSE)
  mod1 <- netsim(est1, param, init, control)

  mod1 <- mutate_epi(mod1, i.prev = i.num / num,
                     i.prev.m2 = i.num.m2 / num.m2)
  expect_equal(names(mod1$epi), c("s.num", "i.num", "num", "s.num.m2", "i.num.m2", "num.m2",
                                  "si.flow", "si.flow.m2", "i.prev", "i.prev.m2"))

  # Add incidence rate per 100 person years (assume time step = 1 week)
  mod1 <- mutate_epi(mod1, ir100 = 5200*(si.flow + si.flow.m2) /
                       (s.num + s.num.m2))
  df <- as.data.frame(mod1)
  expect_true("ir100" %in% names(df))
})

test_that("mutate_epi.dcm", {

  param <- param.dcm(inf.prob = 0.3, inf.prob.g2 = 0.15, balance = "g1")
  init <- init.dcm(s.num = 1000, s.num.g2 = 1000, i.num = 1, i.num.g2 = 0)
  control <- control.dcm(type = "SI", nsteps = 100)
  mod1 <- dcm(param, init, control)

  mod1 <- mutate_epi(mod1, i.prev = i.num / num,
                     i.prev.g2 = i.num.g2 / num.g2)
  expect_equal(names(mod1$epi), c("s.num", "s.num.g2", "i.num", "i.num.g2",
                                  "si.flow", "si.flow.g2", "num", "num.g2",
                                  "i.prev", "i.prev.g2"))

  # Add incidence rate per 100 person years (assume time step = 1 week)
  mod1 <- mutate_epi(mod1, ir100 = 5200*(si.flow + si.flow.g2) /
                       (s.num + s.num.g2))
  df <- as.data.frame(mod1)
  expect_true("ir100" %in% names(df))
})

test_that("mutate_epi.dcm, new mod", {

  syph <- function(t, t0, parms) {
    with(as.list(c(t0, parms)), {

      # 1. track the total population size
      num <- s.num + i.num + r.num

      # 2. define lambda, mu, gamma, and sigma
      ce <- R0 / dur.inf
      lambda <- ce * i.num/num
      mu <- 1/life.expt
      gamma <- 1/dur.inf
      sigma <- 1/dur.imm

      # 3. Write out the four differential equations
      dS <- -lambda*s.num + mu*num - mu*s.num + sigma*r.num
      dI <- lambda*s.num - gamma*i.num - mu*i.num
      dR <- gamma*i.num - mu*r.num - sigma*r.num

      # 4. Outputs
      list(c(dS,
             dI,
             dR))
    })
  }

  param <- param.dcm(R0 = 1.5, life.expt = 365*30,
                     dur.inf = 60, dur.imm = 365*c(8, 10, 12))
  init <- init.dcm(s.num = 1e5, i.num = 1, r.num = 0)
  control <- control.dcm(nsteps = 365, new.mod = syph)
  sim <- dcm(param, init, control)

  sim <- mutate_epi(sim, num = s.num + i.num + r.num)
  sim <- mutate_epi(sim, prev = i.num / num)

  df <- as.data.frame(sim)
  expect_true("num" %in% names(df))
  expect_true("prev" %in% names(df))
  expect_true("num" %in% names(sim$epi))
  expect_true("prev" %in% names(sim$epi))
})

test_that("mutate DCM with constant", {
  param <- param.dcm(inf.prob = 0.2, act.rate = 0.25)
  init <- init.dcm(s.num = 500, i.num = 1)
  control <- control.dcm(type = "SI", nsteps = 500)
  mod1 <- dcm(param, init, control)
  mod1 <- mutate_epi(mod1, prev = i.num/num,
                     cm = 3)
  expect_true(all(sapply(mod1$epi, class) == "data.frame"))
  expect_true(length(unique(sapply(mod1$epi, nrow))) == 1)

  # by itself
  mod1 <- dcm(param, init, control)
  mod1 <- mutate_epi(mod1, cm = 3)
  expect_true(all(sapply(mod1$epi, class) == "data.frame"))
  expect_true(length(unique(sapply(mod1$epi, nrow))) == 1)
})
